#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-06

"""
【节点名称】：
    MOTResultAssembleNode
【依赖项安装】：
    pip install spirems
【订阅类型】：
    spirecv_msgs::2DTargets （目标检测结果）
【发布类型】：
    无
【构造参数说明】：
    parameter_file (str): 全局参数文件（节点参数直接从主节点传过来，所以可以不用传参数文件了）
【节点参数】：
    mot_folder_path (str): MOT测试文件夹地址
    pred_file_path (str): 评价结果文件夹地址
    mot_path (str): MOT根文件夹地址
    mot (str): mot版本
【备注】：
    无
"""

import threading
import os
import subprocess
import cv2
from typing import Union
from queue import Queue
from spirems import Subscriber, Publisher, def_msg, QoS
from spirecv.base.BaseNode import BaseNode
from spirecv.test.mot_metrics.write_val_seqmap import write_val_seqmap


class MOTResultAssembleNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        specified_input_topic: str = None,
        sms_shutdown: bool = True,
        launch_next_emit: bool = False,
        remote_ip: str = '127.0.0.1',
        remote_port: int = 9094,
        mot_folder_path = '',
        pred_file_path = '',
        mot_path = '',
        mot = ''
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown
        )
        self.result_queue = Queue()
        self.queue_pool.append(self.result_queue)

        input_topic_url = '/' + job_name + '/mot/results'if specified_input_topic is None \
            else specified_input_topic
        self._result_reader = Subscriber(
            input_topic_url, 'spirecv_msgs::2DTargets', self.results_callback,
            ip=remote_ip, port=remote_port, qos=QoS.Reliability
        )
        self.launch_next_emit = launch_next_emit
        if self.launch_next_emit:
            self._next_writer = Publisher(
                '/' + job_name + '/launch_next', 'std_msgs::Boolean',
                ip=ip, port=port, qos=QoS.Reliability
            )

        self.file_num = 0
        # self.mot_folder_path = self.get_param("mot_folder_path", os.path.expanduser('~')+"/mot20/train/")
        # self.pred_file_path = self.get_param("pred_file_path", os.path.expanduser('~')+"/mot20/pred_mot20/data/")
        self.mot_folder_path = mot_folder_path
        self.pred_file_path = pred_file_path
        self.mot_path = mot_path
        self.mot = mot
        if not os.path.exists(self.pred_file_path):
            os.makedirs(self.pred_file_path)
        self.mot_files_and_directories = sorted(os.listdir(self.mot_folder_path), key=lambda x: (len(x), x), reverse=False)
        self.end_len = len(os.listdir(os.path.join(self.mot_folder_path,self.mot_files_and_directories[-1],'img1')))
        self.start()

    def release(self):
        BaseNode.release(self)
        self._result_reader.kill()
        if self.launch_next_emit:
            self._next_writer.kill()

    def results_callback(self, msg):
        self.result_queue.put(msg)

    # def file_is_open(self, file_path):
    #     try:
    #         # 尝试打开文件
    #         with open(file_path, 'r') as file:
    #             return False  # 文件没有被打开，返回False
    #     except IOError:
    #         # 如果文件打开失败，抛出IOError异常
    #         return True  # 文件已经被打开，返回True

    def run(self):
        while self.is_running():
            msg = self.result_queue.get(block=True)
            if msg is None:
                break
            # print(msg["img_id"])
            if msg["img_id"]==1:
                self.file_path = self.mot_files_and_directories[self.file_num]
                path = os.path.join(self.pred_file_path,self.file_path+".txt")
                if self.file_num!=0:
                    self.file.close()
                self.file = open(path,"w",encoding='utf-8')
                self.file_num+=1
                print(f'writing {self.file_num}/{len(self.mot_files_and_directories)} files')
            elif msg["img_id"]== self.end_len and self.file_num==len(self.mot_files_and_directories):
                print("write end")
                break
            for i in range(len(msg["targets"])):
                if msg['targets'][i]!=None:
                    self.file.write(f"{msg['img_id']},{msg['targets'][i]['tracked_id']},{msg['targets'][i]['bbox'][0]},{msg['targets'][i]['bbox'][1]},{msg['targets'][i]['bbox'][2]},{msg['targets'][i]['bbox'][3]},0.7,-1,-1,-1\n")
                    #file << frame_id << ","<< target.tracked_id << "," << x << "," << y << "," << width << "," << height << "," << score << "," << "-1,-1,-1" << endl;

        self.file.close()
        
        write_val_seqmap(self.mot_path)
        command = "python3 ../test/mot_metrics/run_mot_challenge.py "+\
            f"--BENCHMARK {self.mot} "+\
                f"--GT_FOLDER {self.mot_folder_path} "+\
                    f"--TRACKERS_FOLDER {os.path.join(self.pred_file_path,'../')} "+\
                        "--SKIP_SPLIT_FOL True "+\
                            f"--SEQMAP_FILE {os.path.join(self.mot_path,'val_seqmap.txt')} "+\
                                "--TRACKERS_TO_EVAL '' "+\
                                    "--METRICS HOTA CLEAR Identity "
        result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        print(result.stdout)
        print(result.stderr)
        cv2.imshow('pedestrian_plot',cv2.imread(os.path.join(self.pred_file_path,'../','pedestrian_plot.png')))
        cv2.waitKey(0)
        cv2.destroyAllWindows()
        

        self.release()
        print('{} quit!'.format(self.__class__.__name__))
    
if __name__ == '__main__':
    node = MOTResultAssembleNode('live')

